import multiprocessing

import pandas as pd
from tqdm import tqdm
import joblib
import collections
import numpy as np


def read_edges_data(data_path):
    '''
    获取连接
    :param data_path: 数据路径
    :return: [[node1,node2]]
    '''
    # 读取图数据
    graph_df = pd.read_csv(data_path, header=None)
    return graph_df.to_numpy()


def read_graph_data(data_path):
    '''
    读取图数据
    :param path: 数据路径
    :return:
    '''

    # 读取图数据
    graph_df = pd.read_csv(data_path, header=None)
    tmp = {}
    # 遍历节点,简化连接数据
    for idx, row in tqdm(graph_df.iterrows(), desc='转换图数据'):
        if row[0] in tmp.keys():
            tmp[row[0]].append(row[1])
        else:
            tmp[row[0]] = [row[1]]

        if row[1] in tmp.keys():
            tmp[row[1]].append(row[0])
        else:
            tmp[row[1]] = [row[0]]

    # 去重处理
    graph = {}
    for k, v in tmp.items():
        graph[k] = list(set(v))
    return len(graph.keys()), graph


def read_graph_embedding(embed_path):
    '''
    读取节点embedding
    :param emb_path: embedding路径
    :return:
    '''
    embed = joblib.load(filename=embed_path)
    result_embed = []
    for k, v in tqdm(embed.items(), desc='获取embdding矩阵'):
        result_embed.append(v)
    return result_embed


def write_embeddings_to_file(generator, discriminator, n_node):
    modes = [generator, discriminator]
    for i in range(2):
        if i == 0:
            path = '../result/generator.emb'
        else:
            path = '../result/discriminator.emb'
        embedding_matrix = modes[i].embedding_matrix
        index = np.array(range(n_node)).reshape(-1, 1)
        embedding_matrix = np.hstack([index, embedding_matrix.detach().cpu().numpy()])
        embedding_list = embedding_matrix.tolist()
        embedding_object = {}
        for emb in embedding_list:
            embedding_object[int(emb[0])] = emb[1:]
        joblib.dump(embedding_object, filename=path)


def construct_trees(nodes, graph):
    """use BFS algorithm to construct the BFS-trees

    Args:
        nodes: the list of nodes in the graph
    Returns:
        trees: dict, root_node_id -> tree, where tree is a dict: node_id -> list: [father, child_0, child_1, ...]
    """

    trees = {}
    for root in tqdm(nodes, desc='生成BFS树'):
        trees[root] = {}
        trees[root][root] = [root]
        used_nodes = set()
        queue = collections.deque([root])
        while len(queue) > 0:
            cur_node = queue.popleft()
            used_nodes.add(cur_node)
            for sub_node in graph[cur_node]:
                if sub_node not in used_nodes:
                    trees[root][cur_node].append(sub_node)
                    trees[root][sub_node] = [cur_node]
                    queue.append(sub_node)
                    used_nodes.add(sub_node)

    return trees

def softmax(x):
    e_x = np.exp(x - np.max(x))  # for computation stability
    return e_x / e_x.sum()
